package com.ndood.core.render;

import java.util.Map;

import org.apache.log4j.Logger;
import org.openqa.selenium.By;
import org.openqa.selenium.Cookie;
import org.openqa.selenium.WebDriver;
import org.openqa.selenium.WebElement;

import us.codecraft.webmagic.Page;
import us.codecraft.webmagic.Request;
import us.codecraft.webmagic.Site;
import us.codecraft.webmagic.Task;
import us.codecraft.webmagic.downloader.Downloader;
import us.codecraft.webmagic.selector.Html;
import us.codecraft.webmagic.selector.PlainText;
/**
 * 自定义Selenium下载器（带渲染功能）
 */
public class SeleniumDownloader implements Downloader {
	Logger logger = Logger.getLogger(getClass());
	
	/**
	 * 下载间隔
	 */
	private int sleepTime = 5000;
	
	/**
	 * 动态渲染，模拟事件
	 */
	private SeleniumDriverAction driverAction;
	
	/**
	 * webDriver池
	 */
	private SeleniumDriverPool driverPool;

	public SeleniumDownloader() {
		super();
	}

	public SeleniumDownloader(int sleepTime, SeleniumDriverPool driverPool) {
		this(sleepTime, driverPool, null);
	}

	public SeleniumDownloader(int sleepTime, SeleniumDriverPool driverPool, SeleniumDriverAction driverAction) {
		this.sleepTime = sleepTime;
		this.driverAction = driverAction;
		if (driverPool != null) {
			this.driverPool = driverPool;
		}
	}

	public SeleniumDownloader setSleepTime(int sleepTime) {
		this.sleepTime = sleepTime;
		return this;
	}

	public void setOperator(SeleniumDriverAction driverAction) {
		this.driverAction = driverAction;
	}
	
	@SuppressWarnings("deprecation")
	@Override
	public Page download(Request request, Task task) {
		// Step1: 从池中获取webDriver
		WebDriver webDriver;
		try {
			webDriver = driverPool.get();
		} catch (Exception e) {
			logger.warn("interrupted", e);
			return null;
		}
		
		// Step2: 下载页面，并休息一段时间
		logger.info("downloading page " + request.getUrl());
		webDriver.get(request.getUrl());
		try {
			Thread.sleep(sleepTime);
		} catch (InterruptedException e) {
			e.printStackTrace();
		}
		
		// Step3: 添加cookie
		WebDriver.Options manage = webDriver.manage();
		manage.window().maximize();
		Site site = task.getSite();
		if (site.getCookies() != null) {
			for (Map.Entry<String, String> cookieEntry : site.getCookies()
					.entrySet()) {
				Cookie cookie = new Cookie(cookieEntry.getKey(),
						cookieEntry.getValue());
				manage.addCookie(cookie);
			}
		}
		
		// Step4: 动态渲染，模拟事件
		if (driverAction != null) {
			try {
				driverAction.execute(webDriver);
			} catch (Exception e) {
				e.printStackTrace();
			}
		}

		// Step5: 封装page
		WebElement webElement = webDriver.findElement(By.xpath("/html"));
		String content = webElement.getAttribute("outerHTML");
		Page page = new Page();
		page.setRawText(content);
		page.setHtml(new Html(content, request.getUrl()));
		page.setUrl(new PlainText(request.getUrl()));
		page.setRequest(request);
		driverPool.returnToPool(webDriver);
		return page;
	}

	@Override
	public void setThread(int thread) {}
	
}
